# gomoku_ai/main.py
import argparse
import os
import sys
import time

# 导入配置
from config import *

# --- 从 device_config 导入设备 ---
from utils.device_config import DEVICE

# 导入模块
from game.gomoku_game import GomokuGame
from model.gomoku_net import GomokuNet
from trainer.trainer import Trainer
from utils.play import play_vs_human


def find_latest_checkpoint(checkpoint_dir, template):
    """找到最新的检查点文件"""
    if not os.path.exists(checkpoint_dir):
        return None

    checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith("gomoku_model_iter_") and f.endswith(".pth")]
    if not checkpoints:
        return None

    checkpoints.sort(key=lambda x: int(x.split("_iter_")[1].split(".pth")[0]))
    return os.path.join(checkpoint_dir, checkpoints[-1])


def train(visual_mode):
    """训练模型"""
    print("--- 开始训练 ---")

    # 创建游戏和神经网络
    game = GomokuGame(board_size=BOARD_SIZE, visual_mode=visual_mode)
    net = GomokuNet(game, learning_rate=LEARNING_RATE, device=DEVICE)

    # 加载最新模型
    latest_model_path = find_latest_checkpoint(CHECKPOINT_DIR, MODEL_FILENAME_TEMPLATE)
    if latest_model_path:
        net.load(latest_model_path)
        print(f"✅ 已加载模型: {latest_model_path}")
    else:
        print("⚠️ 未找到已有模型，从头开始训练。")

    # 创建训练器
    trainer = Trainer(game, net, mcts_simulations=MCTS_SIMULATIONS, checkpoint_dir=CHECKPOINT_DIR)

    # 开始训练
    start_time = time.time()
    trainer.learn(
        num_iterations=NUM_ITERATIONS,
        num_episodes_per_iter=NUM_EPISODES_PER_ITER,
        batch_size=BATCH_SIZE,
        epochs_per_iter=TRAIN_EPOCHS_PER_ITER,
        model_filename_template=MODEL_FILENAME_TEMPLATE
    )
    end_time = time.time()

    print(f"🎉 训练完成！耗时: {end_time - start_time:.2f}秒")

    # 关闭图形界面
    if game.visual_board:
        game.visual_board.close()


def play():
    """与AI对弈"""
    print("--- 开始与AI对弈 ---")

    # 创建游戏和神经网络
    game = GomokuGame(board_size=BOARD_SIZE, visual_mode=True)
    net = GomokuNet(game, learning_rate=LEARNING_RATE, device=DEVICE)

    # 加载最新模型
    latest_model_path = find_latest_checkpoint(CHECKPOINT_DIR, MODEL_FILENAME_TEMPLATE)
    if not latest_model_path:
        print("❌ 未找到训练好的模型。请先运行 'python main.py train' 来训练模型。")
        return

    net.load(latest_model_path)
    print(f"✅ 已加载模型: {latest_model_path}")

    # 开始对弈
    play_vs_human(game, net, mcts_simulations=MCTS_SIMULATIONS * 2)

    # 关闭图形界面
    if game.visual_board:
        game.visual_board.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="五子棋 AI 训练与对战系统")
    parser.add_argument("command", choices=["train", "play"], help="命令: 'train' 训练模型, 'play' 与AI对弈")
    parser.add_argument("--no-visual", action="store_true", help="强制训练时不显示图形界面（会覆盖config.py中的设置）")

    args = parser.parse_args()

    if args.command == "train":
        # --- 关键修改：新的逻辑 ---
        # visual_mode 的值由 VISUAL_TRAINING 和 --no-visual 共同决定
        visual_mode = VISUAL_TRAINING and not args.no_visual
        train(visual_mode)
    elif args.command == "play":
        play()
